# from pytypes import typechecked
from common.clickhouse_alc_engine import get_db_engine
from datetime import datetime, timedelta
import numpy as np
import pandas as pd
from scipy.special import comb
from sqlalchemy import text
import talib
from typing import Literal


def get_data_from_clickhouse(
    symbol: str, start: str, end: str, timeframe: str
) -> pd.DataFrame:
    engine = get_db_engine("coins")
    start = datetime.strptime(start, "%Y%m%d")
    end = datetime.strptime(end, "%Y%m%d") + timedelta(days=1)
    table = f"klines_{timeframe}"
    with engine.connect() as connection:
        result = connection.execute(
            text(
                "select time, symbol, open, high, low, close, volume from :tb where symbol=:symbol and time >= :start and time < :end "
            ),
            {"tb": table, "start": start, "end": end, "symbol": symbol},
        )
        klines = result.fetchall()
    klines = pd.DataFrame(klines)
    klines.set_index("time", inplace=True)
    klines.sort_index(inplace=True)
    return klines


def cal_rvi_factor_from_klines(df: pd.DataFrame, N: int = 14) -> pd.DataFrame:
    a = df["close"] - df["open"]
    e = df["high"] - df["low"]
    NUMERATOR = (a + 2 * a.shift(1) + 2 * a.shift(2) + a.shift(3)) / 6
    DENOMINATOR = (e + 2 * e.shift(1) + 2 * e.shift(2) + e.shift(3)) / 6
    sma_NUMERATOR = talib.SMA(NUMERATOR, timeperiods=N)
    sma_DENOMINATOR = talib.SMA(DENOMINATOR, timeperiods=N)
    r = sma_NUMERATOR / sma_DENOMINATOR
    value = (r + 2 * r.shift(1) + 2 * r.shift(2) + r.shift(3)) / 6
    return value


def cal_pascal_factor(df: pd.DataFrame, N: int = 14) -> pd.Series:
    weights = np.array([comb(N, k) for k in range(N)])
    result = np.convolve(df["close"], weights, mode="valid")
    return pd.Series(result, index=df.index)


def get_cfo(prices: pd.Series):
    """使用线性回归计算 FMA"""
    X = np.arange(len(prices))
    y = prices
    coes = np.polyfit(X, y)
    forecast = np.polyval(coes, [len(prices)])
    return forecast


def cal_cfo_factor(
    df: pd.DataFrame,
    N: int = 14,
    kind: Literal["O", "H", "L", "C", "HL", "HLC", "HLCC", "OHLC"] = "C",
) -> pd.Series:
    dc = {"O": df["open"], "C": df["close"], "H": df["high"], "L": df["low"]}
    value = dc[kind[0]]
    for x in dc[1:]:
        value = value + dc[x]
    value: pd.Series = value / len(kind)
    result = value.rolling(window=N, center=False).apply(get_cfo)
    return result


def add_to_factor(symbol: str, fname: str, start: str, end: str, timeframe: str):
    start = datetime.strptime(start, "%Y%m%d")
    end = datetime.strptime(end, "%Y%m%d") + timedelta(days=1)

    table = f"factor_{timeframe}"
    timeframe = timeframe.replace("m", "T")
    indices = pd.date_range(start=start, end=end, freq=timeframe)[:-1]
    df = pd.DataFrame(index=indices, columns=["symbol", "fname", "value"])
    df["symbol"] = symbol
    df["fname"] = fname
    df["value"] = np.random.random(len(df)) * 2 - 1
    df["time"] = indices
    df.to_sql(table, get_db_engine(), index=False, if_exists="append")


if __name__ == "__main__":
    add_to_factor("BTC/USDT", "rand2", "20200101", "20250208", "1m")
